package main

import (
	"archive/tar"
	"compress/gzip"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
)

// CompressDir 压缩整个目录到指定的tar.gz文件
func CompressDir(srcDir, dstFile string) error {
	// 创建输出文件
	out, err := os.Create(dstFile)
	if err != nil {
		return fmt.Errorf("创建输出文件失败: %w", err)
	}
	defer out.Close()

	// 创建gzip写入器
	gzw := gzip.NewWriter(out)
	defer gzw.Close()

	// 创建tar写入器
	tw := tar.NewWriter(gzw)
	defer tw.Close()

	// 遍历源目录
	return filepath.Walk(srcDir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		// 创建tar头
		header, err := tar.FileInfoHeader(info, info.Name())
		if err != nil {
			return fmt.Errorf("创建tar头失败: %w", err)
		}

		// 处理文件路径（去除源目录前缀）
		header.Name = strings.TrimPrefix(path, filepath.Dir(srcDir)+string(os.PathSeparator))

		// 写入tar头
		if err := tw.WriteHeader(header); err != nil {
			return fmt.Errorf("写入tar头失败: %w", err)
		}

		// 如果是目录，不需要进一步处理
		if info.IsDir() {
			return nil
		}

		// 打开文件
		file, err := os.Open(path)
		if err != nil {
			return fmt.Errorf("打开文件失败: %w", err)
		}
		defer file.Close()

		// 复制文件内容到tar写入器
		_, err = io.Copy(tw, file)
		if err != nil {
			return fmt.Errorf("复制文件内容失败: %w", err)
		}

		return nil
	})
}

func main() {
	srcDir := "./testdir"       // 要压缩的源目录
	dstFile := "testdir.tar.gz" // 输出的压缩文件名

	if err := CompressDir(srcDir, dstFile); err != nil {
		fmt.Println("压缩失败:", err)
		return
	}

	fmt.Printf("目录 %s 已成功压缩到 %s\n", srcDir, dstFile)
}
